import numpy as np

class KMeans:
     def __init__(self,data, num_clusters ):
          self.data = data
          self.num_clusters = num_clusters

     @staticmethod
     def centroid_init(data, num_clusters):
          num_examples = data.shape[0]
          random_ids = np.random.permutation(num_examples) # 用于随机算法进行重新洗牌
          centroids = data[random_ids[:num_clusters], :] #随机生成质心
          return centroids 
     @staticmethod
     def centroid_find_closest(data, centroids):
          # num_examples  样本数
          num_examples = data.shape[0]
          # 质心的个数
          num_centroids = centroids.shape[0]
          closest_centroids_ids = np.zeros((num_examples,1))
          for example_index in range(num_examples):
               distance = np.zeros((num_centroids, 1))
               for centroid_index in range(num_centroids):
                    distance_diff = data[example_index, :] - centroids[centroid_index, :]
                    distance[centroid_index] = np.sum(distance_diff ** 2)
               closest_centroids_ids[example_index] = np.argmin(distance)
          return closest_centroids_ids
     
     @staticmethod
     def centroids_compute(data, closest_centroids_ids,  num_clusters):
          # 计算均值
          # 计算有多少个features
          num_features = data.shape[1]
          # number of clusters [ k, num_features]
          centroids = np.zeros((num_clusters, num_features))
          for centroid_id in range(num_clusters):
               closest_ids = closest_centroids_ids == centroid_id
               centroids[closest_ids] = np.mean(data[closest_ids.flatten(), :], axis=0)
          return centroids


     def train(self, max_iterations):
          # 1. 先随机生成k个质心
          centroids = KMeans.centroid_init(self.data, self.num_clusters)
          # 2. 计算样本中的所有点到中心点的距离，每个点计算三个距离
          # start training
          num_examples = self.data.shape[0]
          closest_centroids_ids = np.empty((num_examples, 1)) #  寻找最近的中心点
          for _ in range(max_iterations):
               # 3.  得到current样本点到中心点的距离
               closest_centroids_ids = KMeans.centroid_find_closest(self.data, centroids)
               # 4. update centroid 进行中心点位置更新
               centroids = KMeans. centroids_compute(self.data, closest_centroids_ids, self.num_clusters)
          return centroids, closest_centroids_ids #返回了中心点和距离
